use rusqlite::{Connection, Result as SqliteResult, OptionalExtension};
use serde::{Deserialize, Serialize};
use std::path::Path;
use std::sync::{Arc, Mutex};
use uuid::Uuid;
use anyhow::Result;
use tracing;

use crate::vector_search::{HybridSearchManager, SearchResult};

// Database models
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Conversation {
    pub id: i64,
    pub uuid: String,
    pub title: String,
    pub agent_id: Option<i64>,
    pub created_at: String,
    pub updated_at: String,
    pub message_count: i64,
    pub last_message: Option<String>,
    pub metadata: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    pub id: i64,
    pub uuid: String,
    pub conversation_id: i64,
    pub role: String,
    pub content: String,
    pub created_at: String,
    pub status: String,
    pub metadata: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Bookmark {
    pub id: i64,
    pub uuid: String,
    pub title: String,
    pub prompt: String,
    pub description: Option<String>,
    pub category: Option<String>,
    pub is_favorite: bool,
    pub usage_count: i64,
    pub created_at: String,
    pub updated_at: String,
    pub tags: Option<String>, // JSON array
    pub metadata: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KnowledgeItem {
    pub id: i64,
    pub uuid: String,
    pub title: String,
    pub content: Option<String>,
    pub item_type: String, // "document", "url", "text"
    pub url: Option<String>,
    pub file_path: Option<String>,
    pub file_size: Option<i64>,
    pub mime_type: Option<String>,
    pub created_at: String,
    pub updated_at: String,
    pub tags: Option<String>, // JSON string
    pub metadata: Option<String>, // JSON string
}

impl Default for KnowledgeItem {
    fn default() -> Self {
        Self {
            id: 0,
            uuid: String::new(),
            title: String::new(),
            content: None,
            item_type: String::new(),
            url: None,
            file_path: None,
            file_size: None,
            mime_type: None,
            created_at: String::new(),
            updated_at: String::new(),
            tags: None,
            metadata: None,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KnowledgeCollection {
    pub id: i64,
    pub uuid: String,
    pub name: String,
    pub description: Option<String>,
    pub created_at: String,
    pub updated_at: String,
    pub metadata: Option<String>,
}

// Database manager
#[derive(Debug, Clone)]
pub struct DatabaseManager {
    connection: Arc<Mutex<Connection>>,
    vector_search: Option<Arc<HybridSearchManager>>,
}

impl DatabaseManager {
    pub fn new<P: AsRef<Path>>(db_path: P) -> SqliteResult<Self> {
        let conn = Connection::open(db_path)?;

        // Run migrations
        Self::run_migrations(&conn)?;

        Ok(Self {
            connection: Arc::new(Mutex::new(conn)),
            vector_search: None,
        })
    }

    pub async fn new_with_vector_search<P: AsRef<Path>>(db_path: P) -> Result<Self> {
        let path_str = db_path.as_ref().to_str()
            .ok_or_else(|| anyhow::anyhow!("Invalid database path"))?;
        
        let conn = Connection::open(&db_path)?;

        // Run migrations
        Self::run_migrations(&conn)?;

        // Initialize vector search
        let vector_search = HybridSearchManager::new(path_str).await?;

        Ok(Self {
            connection: Arc::new(Mutex::new(conn)),
            vector_search: Some(Arc::new(vector_search)),
        })
    }

    pub async fn new_with_vector_search_and_agent_manager<P: AsRef<Path>>(
        db_path: P, 
        agent_manager: rig_app::agents::AgentManager
    ) -> Result<Self> {
        let path_str = db_path.as_ref().to_str()
            .ok_or_else(|| anyhow::anyhow!("Invalid database path"))?;
        
        let conn = Connection::open(&db_path)?;

        // Run migrations
        Self::run_migrations(&conn)?;

        // Initialize vector search with agent manager
        let vector_search = HybridSearchManager::new_with_agent_manager(path_str, agent_manager).await?;

        Ok(Self {
            connection: Arc::new(Mutex::new(conn)),
            vector_search: Some(Arc::new(vector_search)),
        })
    }

    fn run_migrations(conn: &Connection) -> SqliteResult<()> {
        // Read and execute the migration SQL
        let migration_sql = include_str!("../migrations/001_initial_schema.sql");
        conn.execute_batch(migration_sql)?;
        
        // Run vector search migration
        let vector_migration_sql = include_str!("../migrations/002_vector_search.sql");
        conn.execute_batch(vector_migration_sql)?;
        
        Ok(())
    }

    // Conversation methods
    pub fn create_conversation(
        &self,
        title: &str,
        agent_id: Option<i64>,
        metadata: Option<&str>,
    ) -> SqliteResult<i64> {
        let conn = self.connection.lock().unwrap();
        let uuid = Uuid::new_v4().to_string();

        conn.execute(
            "INSERT INTO conversations (uuid, title, agent_id, metadata) VALUES (?1, ?2, ?3, ?4)",
            (uuid, title, agent_id, metadata),
        )?;

        Ok(conn.last_insert_rowid())
    }

    pub fn get_conversation(&self, id: i64) -> SqliteResult<Option<Conversation>> {
        let conn = self.connection.lock().unwrap();
        let mut stmt = conn.prepare(
            "SELECT id, uuid, title, agent_id, created_at, updated_at, message_count, last_message, metadata 
             FROM conversations WHERE id = ?1"
        )?;

        let conversation = stmt
            .query_row([id], |row| {
                Ok(Conversation {
                    id: row.get(0)?,
                    uuid: row.get(1)?,
                    title: row.get(2)?,
                    agent_id: row.get(3)?,
                    created_at: row.get(4)?,
                    updated_at: row.get(5)?,
                    message_count: row.get(6)?,
                    last_message: row.get(7)?,
                    metadata: row.get(8)?,
                })
            })
            .optional()?;

        Ok(conversation)
    }

    pub fn list_conversations(
        &self,
        limit: Option<i64>,
        offset: Option<i64>,
    ) -> SqliteResult<Vec<Conversation>> {
        let conn = self.connection.lock().unwrap();
        let limit = limit.unwrap_or(50);
        let offset = offset.unwrap_or(0);

        let mut stmt = conn.prepare(
            "SELECT id, uuid, title, agent_id, created_at, updated_at, message_count, last_message, metadata 
             FROM conversations ORDER BY updated_at DESC LIMIT ?1 OFFSET ?2"
        )?;

        let conversations = stmt
            .query_map([limit, offset], |row| {
                Ok(Conversation {
                    id: row.get(0)?,
                    uuid: row.get(1)?,
                    title: row.get(2)?,
                    agent_id: row.get(3)?,
                    created_at: row.get(4)?,
                    updated_at: row.get(5)?,
                    message_count: row.get(6)?,
                    last_message: row.get(7)?,
                    metadata: row.get(8)?,
                })
            })?
            .collect::<SqliteResult<Vec<_>>>()?;

        Ok(conversations)
    }

    pub fn update_conversation(
        &self,
        id: i64,
        title: Option<&str>,
        metadata: Option<&str>,
    ) -> SqliteResult<bool> {
        let conn = self.connection.lock().unwrap();

        let mut query = String::from("UPDATE conversations SET ");
        let mut params = Vec::new();
        let mut updates = Vec::new();

        if let Some(title) = title {
            updates.push("title = ?");
            params.push(title);
        }

        if let Some(metadata) = metadata {
            updates.push("metadata = ?");
            params.push(metadata);
        }

        if updates.is_empty() {
            return Ok(false);
        }

        query.push_str(&updates.join(", "));
        query.push_str(" WHERE id = ?");
        let id_str = id.to_string();
        params.push(&id_str);

        let rows_affected = conn.execute(&query, rusqlite::params_from_iter(params))?;
        Ok(rows_affected > 0)
    }

    pub fn delete_conversation(&self, id: i64) -> SqliteResult<bool> {
        let conn = self.connection.lock().unwrap();
        let rows_affected = conn.execute("DELETE FROM conversations WHERE id = ?1", [id])?;
        Ok(rows_affected > 0)
    }

    // Message methods
    pub fn create_message(
        &self,
        conversation_id: i64,
        role: &str,
        content: &str,
        status: Option<&str>,
        metadata: Option<&str>,
    ) -> SqliteResult<i64> {
        let conn = self.connection.lock().unwrap();
        let uuid = Uuid::new_v4().to_string();
        let status = status.unwrap_or("sent");

        // Start transaction
        let tx = conn.unchecked_transaction()?;

        // Insert message
        tx.execute(
            "INSERT INTO messages (uuid, conversation_id, role, content, status, metadata) 
             VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
            (uuid, conversation_id, role, content, status, metadata),
        )?;

        let message_id = tx.last_insert_rowid();

        // Update conversation message count and last message
        tx.execute(
            "UPDATE conversations 
             SET message_count = message_count + 1, last_message = ?1 
             WHERE id = ?2",
            (content, conversation_id),
        )?;

        tx.commit()?;
        Ok(message_id)
    }

    pub fn get_messages(
        &self,
        conversation_id: i64,
        limit: Option<i64>,
        offset: Option<i64>,
    ) -> SqliteResult<Vec<Message>> {
        let conn = self.connection.lock().unwrap();
        let limit = limit.unwrap_or(100);
        let offset = offset.unwrap_or(0);

        let mut stmt = conn.prepare(
            "SELECT id, uuid, conversation_id, role, content, created_at, status, metadata 
             FROM messages WHERE conversation_id = ?1 ORDER BY created_at ASC LIMIT ?2 OFFSET ?3",
        )?;

        let messages = stmt
            .query_map([conversation_id, limit, offset], |row| {
                Ok(Message {
                    id: row.get(0)?,
                    uuid: row.get(1)?,
                    conversation_id: row.get(2)?,
                    role: row.get(3)?,
                    content: row.get(4)?,
                    created_at: row.get(5)?,
                    status: row.get(6)?,
                    metadata: row.get(7)?,
                })
            })?
            .collect::<SqliteResult<Vec<_>>>()?;

        Ok(messages)
    }

    pub fn update_message(
        &self,
        id: i64,
        content: Option<&str>,
        status: Option<&str>,
        metadata: Option<&str>,
    ) -> SqliteResult<bool> {
        let conn = self.connection.lock().unwrap();

        let mut query = String::from("UPDATE messages SET ");
        let mut params = Vec::new();
        let mut updates = Vec::new();

        if let Some(content) = content {
            updates.push("content = ?");
            params.push(content);
        }

        if let Some(status) = status {
            updates.push("status = ?");
            params.push(status);
        }

        if let Some(metadata) = metadata {
            updates.push("metadata = ?");
            params.push(metadata);
        }

        if updates.is_empty() {
            return Ok(false);
        }

        query.push_str(&updates.join(", "));
        query.push_str(" WHERE id = ?");
        let id_str2 = id.to_string();
        params.push(&id_str2);

        let rows_affected = conn.execute(&query, rusqlite::params_from_iter(params))?;
        Ok(rows_affected > 0)
    }

    // Bookmark methods
    pub fn create_bookmark(
        &self,
        title: &str,
        prompt: &str,
        description: Option<&str>,
        category: Option<&str>,
        tags: Option<&str>,
        metadata: Option<&str>,
    ) -> SqliteResult<i64> {
        let conn = self.connection.lock().unwrap();
        let uuid = Uuid::new_v4().to_string();

        conn.execute(
            "INSERT INTO bookmarks (uuid, title, prompt, description, category, tags, metadata) 
             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
            (uuid, title, prompt, description, category, tags, metadata),
        )?;

        Ok(conn.last_insert_rowid())
    }

    pub fn get_bookmark(&self, id: i64) -> SqliteResult<Option<Bookmark>> {
        let conn = self.connection.lock().unwrap();
        let mut stmt = conn.prepare(
            "SELECT id, uuid, title, prompt, description, category, is_favorite, usage_count, 
                    created_at, updated_at, tags, metadata 
             FROM bookmarks WHERE id = ?1",
        )?;

        let bookmark = stmt
            .query_row([id], |row| {
                Ok(Bookmark {
                    id: row.get(0)?,
                    uuid: row.get(1)?,
                    title: row.get(2)?,
                    prompt: row.get(3)?,
                    description: row.get(4)?,
                    category: row.get(5)?,
                    is_favorite: row.get(6)?,
                    usage_count: row.get(7)?,
                    created_at: row.get(8)?,
                    updated_at: row.get(9)?,
                    tags: row.get(10)?,
                    metadata: row.get(11)?,
                })
            })
            .optional()?;

        Ok(bookmark)
    }

    pub fn list_bookmarks(
        &self,
        category: Option<&str>,
        favorites_only: bool,
        limit: Option<i64>,
        offset: Option<i64>,
    ) -> SqliteResult<Vec<Bookmark>> {
        let conn = self.connection.lock().unwrap();
        let limit = limit.unwrap_or(50);
        let offset = offset.unwrap_or(0);

        let mut query = String::from(
            "SELECT id, uuid, title, prompt, description, category, is_favorite, usage_count, 
                    created_at, updated_at, tags, metadata 
             FROM bookmarks WHERE 1=1",
        );
        let mut params = Vec::new();

        if let Some(category) = category {
            query.push_str(" AND category = ?");
            params.push(category);
        }

        if favorites_only {
            query.push_str(" AND is_favorite = 1");
        }

        query.push_str(" ORDER BY created_at DESC LIMIT ? OFFSET ?");
        let limit_str = limit.to_string();
        let offset_str = offset.to_string();
        params.push(&limit_str);
        params.push(&offset_str);

        let mut stmt = conn.prepare(&query)?;
        let bookmarks = stmt
            .query_map(rusqlite::params_from_iter(params), |row| {
                Ok(Bookmark {
                    id: row.get(0)?,
                    uuid: row.get(1)?,
                    title: row.get(2)?,
                    prompt: row.get(3)?,
                    description: row.get(4)?,
                    category: row.get(5)?,
                    is_favorite: row.get(6)?,
                    usage_count: row.get(7)?,
                    created_at: row.get(8)?,
                    updated_at: row.get(9)?,
                    tags: row.get(10)?,
                    metadata: row.get(11)?,
                })
            })?
            .collect::<SqliteResult<Vec<_>>>()?;

        Ok(bookmarks)
    }

    pub fn update_bookmark_usage(&self, id: i64) -> SqliteResult<bool> {
        let conn = self.connection.lock().unwrap();
        let rows_affected = conn.execute(
            "UPDATE bookmarks SET usage_count = usage_count + 1 WHERE id = ?1",
            [id],
        )?;
        Ok(rows_affected > 0)
    }

    pub fn toggle_bookmark_favorite(&self, id: i64) -> SqliteResult<bool> {
        let conn = self.connection.lock().unwrap();
        let rows_affected = conn.execute(
            "UPDATE bookmarks SET is_favorite = NOT is_favorite WHERE id = ?1",
            [id],
        )?;
        Ok(rows_affected > 0)
    }

    pub fn delete_bookmark(&self, id: i64) -> SqliteResult<bool> {
        let conn = self.connection.lock().unwrap();
        let rows_affected = conn.execute("DELETE FROM bookmarks WHERE id = ?1", [id])?;
        Ok(rows_affected > 0)
    }

    // Knowledge item methods
    pub fn create_knowledge_item(
        &self,
        title: &str,
        content: Option<&str>,
        item_type: &str,
        url: Option<&str>,
        file_path: Option<&str>,
        file_size: Option<i64>,
        mime_type: Option<&str>,
        tags: Option<&str>,
        metadata: Option<&str>,
    ) -> SqliteResult<i64> {
        let conn = self.connection.lock().unwrap();
        let uuid = Uuid::new_v4().to_string();

        conn.execute(
            "INSERT INTO knowledge_items (uuid, title, content, item_type, url, file_path, file_size, mime_type, tags, metadata) 
             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
            (uuid, title, content, item_type, url, file_path, file_size, mime_type, tags, metadata),
        )?;

        Ok(conn.last_insert_rowid())
    }

    pub fn get_knowledge_item(&self, id: i64) -> SqliteResult<Option<KnowledgeItem>> {
        let conn = self.connection.lock().unwrap();
        let mut stmt = conn.prepare(
            "SELECT id, uuid, title, content, item_type, url, file_path, file_size, mime_type, 
                    created_at, updated_at, tags, metadata 
             FROM knowledge_items WHERE id = ?1",
        )?;

        let item = stmt
            .query_row([id], |row| {
                Ok(KnowledgeItem {
                    id: row.get(0)?,
                    uuid: row.get(1)?,
                    title: row.get(2)?,
                    content: row.get(3)?,
                    item_type: row.get(4)?,
                    url: row.get(5)?,
                    file_path: row.get(6)?,
                    file_size: row.get(7)?,
                    mime_type: row.get(8)?,
                    created_at: row.get(9)?,
                    updated_at: row.get(10)?,
                    tags: row.get(11)?,
                    metadata: row.get(12)?,
                })
            })
            .optional()?;

        Ok(item)
    }

    pub fn list_knowledge_items(
        &self,
        item_type: Option<&str>,
        limit: Option<i64>,
        offset: Option<i64>,
    ) -> SqliteResult<Vec<KnowledgeItem>> {
        let conn = self.connection.lock().unwrap();
        let limit = limit.unwrap_or(50);
        let offset = offset.unwrap_or(0);

        let mut query = String::from(
            "SELECT id, uuid, title, content, item_type, url, file_path, file_size, mime_type, 
                    created_at, updated_at, tags, metadata 
             FROM knowledge_items WHERE 1=1",
        );
        let mut params = Vec::new();

        if let Some(item_type) = item_type {
            query.push_str(" AND item_type = ?");
            params.push(item_type);
        }

        query.push_str(" ORDER BY created_at DESC LIMIT ? OFFSET ?");
        let limit_str2 = limit.to_string();
        let offset_str2 = offset.to_string();
        params.push(&limit_str2);
        params.push(&offset_str2);

        let mut stmt = conn.prepare(&query)?;
        let items = stmt
            .query_map(rusqlite::params_from_iter(params), |row| {
                Ok(KnowledgeItem {
                    id: row.get(0)?,
                    uuid: row.get(1)?,
                    title: row.get(2)?,
                    content: row.get(3)?,
                    item_type: row.get(4)?,
                    url: row.get(5)?,
                    file_path: row.get(6)?,
                    file_size: row.get(7)?,
                    mime_type: row.get(8)?,
                    created_at: row.get(9)?,
                    updated_at: row.get(10)?,
                    tags: row.get(11)?,
                    metadata: row.get(12)?,
                })
            })?
            .collect::<SqliteResult<Vec<_>>>()?;

        Ok(items)
    }

    pub fn delete_knowledge_item(&self, id: i64) -> SqliteResult<bool> {
        let conn = self.connection.lock().unwrap();
        let rows_affected = conn.execute("DELETE FROM knowledge_items WHERE id = ?1", [id])?;
        Ok(rows_affected > 0)
    }

    // Search methods
    pub fn search_bookmarks(&self, query: &str, limit: Option<i64>) -> SqliteResult<Vec<Bookmark>> {
        let conn = self.connection.lock().unwrap();
        let limit = limit.unwrap_or(20);
        let search_query = format!("%{}%", query);

        let mut stmt = conn.prepare(
            "SELECT id, uuid, title, prompt, description, category, is_favorite, usage_count, 
                    created_at, updated_at, tags, metadata 
             FROM bookmarks 
             WHERE title LIKE ?1 OR prompt LIKE ?1 OR description LIKE ?1 
             ORDER BY usage_count DESC, created_at DESC 
             LIMIT ?2",
        )?;

        let bookmarks = stmt
            .query_map([&search_query, &limit.to_string()], |row| {
                Ok(Bookmark {
                    id: row.get(0)?,
                    uuid: row.get(1)?,
                    title: row.get(2)?,
                    prompt: row.get(3)?,
                    description: row.get(4)?,
                    category: row.get(5)?,
                    is_favorite: row.get(6)?,
                    usage_count: row.get(7)?,
                    created_at: row.get(8)?,
                    updated_at: row.get(9)?,
                    tags: row.get(10)?,
                    metadata: row.get(11)?,
                })
            })?
            .collect::<SqliteResult<Vec<_>>>()?;

        Ok(bookmarks)
    }

    pub fn search_knowledge_items(
        &self,
        query: &str,
        limit: Option<i64>,
    ) -> SqliteResult<Vec<KnowledgeItem>> {
        let conn = self.connection.lock().unwrap();
        let limit = limit.unwrap_or(20);
        let search_query = format!("%{}%", query);

        let mut stmt = conn.prepare(
            "SELECT id, uuid, title, content, item_type, url, file_path, file_size, mime_type, 
                    created_at, updated_at, tags, metadata 
             FROM knowledge_items 
             WHERE title LIKE ?1 OR content LIKE ?1 
             ORDER BY created_at DESC 
             LIMIT ?2",
        )?;

        let items = stmt
            .query_map([&search_query, &limit.to_string()], |row| {
                Ok(KnowledgeItem {
                    id: row.get(0)?,
                    uuid: row.get(1)?,
                    title: row.get(2)?,
                    content: row.get(3)?,
                    item_type: row.get(4)?,
                    url: row.get(5)?,
                    file_path: row.get(6)?,
                    file_size: row.get(7)?,
                    mime_type: row.get(8)?,
                    created_at: row.get(9)?,
                    updated_at: row.get(10)?,
                    tags: row.get(11)?,
                    metadata: row.get(12)?,
                })
            })?
            .collect::<SqliteResult<Vec<_>>>()?;

        Ok(items)
    }

    // Vector search methods
    pub async fn vector_search_knowledge_items(
        &self,
        query: &str,
        limit: Option<usize>,
    ) -> Result<Vec<SearchResult>> {
        self.vector_search_knowledge_items_with_agent(query, limit, None).await
    }

    pub async fn vector_search_knowledge_items_with_agent(
        &self,
        query: &str,
        limit: Option<usize>,
        agent_id: Option<i64>,
    ) -> Result<Vec<SearchResult>> {
        if let Some(vector_search) = &self.vector_search {
            let limit = limit.unwrap_or(10);
            vector_search.hybrid_search_with_agent(query, limit, 0.7, agent_id).await
        } else {
            Err(anyhow::anyhow!("Vector search not initialized"))
        }
    }

    pub fn search_knowledge_items_fts(
        &self,
        query: &str,
        limit: Option<i64>,
    ) -> SqliteResult<Vec<KnowledgeItem>> {
        let conn = self.connection.lock().unwrap();
        let limit = limit.unwrap_or(20);

        let mut stmt = conn.prepare(
            "SELECT ki.id, ki.uuid, ki.title, ki.content, ki.item_type, ki.url, ki.file_path, 
                    ki.file_size, ki.mime_type, ki.created_at, ki.updated_at, ki.tags, ki.metadata 
             FROM knowledge_items ki
             JOIN knowledge_items_fts fts ON ki.id = fts.rowid
             WHERE knowledge_items_fts MATCH ?1
             ORDER BY rank
             LIMIT ?2",
        )?;

        let items = stmt
            .query_map([query, &limit.to_string()], |row| {
                Ok(KnowledgeItem {
                    id: row.get(0)?,
                    uuid: row.get(1)?,
                    title: row.get(2)?,
                    content: row.get(3)?,
                    item_type: row.get(4)?,
                    url: row.get(5)?,
                    file_path: row.get(6)?,
                    file_size: row.get(7)?,
                    mime_type: row.get(8)?,
                    created_at: row.get(9)?,
                    updated_at: row.get(10)?,
                    tags: row.get(11)?,
                    metadata: row.get(12)?,
                })
            })?
            .collect::<SqliteResult<Vec<_>>>()?;

        Ok(items)
    }

    pub async fn add_knowledge_item_with_vector(
        &self,
        title: &str,
        content: Option<&str>,
        item_type: &str,
        url: Option<&str>,
        file_path: Option<&str>,
        file_size: Option<i64>,
        mime_type: Option<&str>,
        tags: Option<&str>,
        metadata: Option<&str>,
    ) -> Result<i64> {
        self.add_knowledge_item_with_vector_and_agent(
            title, content, item_type, url, file_path, file_size, mime_type, tags, metadata, None
        ).await
    }

    pub async fn add_knowledge_item_with_vector_and_agent(
        &self,
        title: &str,
        content: Option<&str>,
        item_type: &str,
        url: Option<&str>,
        file_path: Option<&str>,
        file_size: Option<i64>,
        mime_type: Option<&str>,
        tags: Option<&str>,
        metadata: Option<&str>,
        agent_id: Option<i64>,
    ) -> Result<i64> {
        // Create the knowledge item first
        let item_id = self.create_knowledge_item(
            title, content, item_type, url, file_path, file_size, mime_type, tags, metadata,
        )?;

        // Get the created item and add to vector store
        if let Ok(Some(item)) = self.get_knowledge_item(item_id) {
            if let Some(vector_search) = &self.vector_search {
                if let Err(e) = vector_search.add_knowledge_item_with_agent(&item, agent_id).await {
                    tracing::warn!("Failed to add item to vector store: {}", e);
                }
            }
        }

        Ok(item_id)
    }

    pub async fn update_knowledge_item_with_vector(
        &self,
        id: i64,
        title: Option<&str>,
        content: Option<&str>,
        tags: Option<&str>,
        metadata: Option<&str>,
    ) -> Result<bool, anyhow::Error> {
        self.update_knowledge_item_with_vector_and_agent(id, title, content, tags, metadata, None).await
    }

    pub async fn update_knowledge_item_with_vector_and_agent(
        &self,
        id: i64,
        title: Option<&str>,
        content: Option<&str>,
        tags: Option<&str>,
        metadata: Option<&str>,
        agent_id: Option<i64>,
    ) -> Result<bool, anyhow::Error> {
        let title = title.map(|s| s.to_string());
        let content = content.map(|s| s.to_string());
        let tags = tags.map(|s| s.to_string());
        let metadata = metadata.map(|s| s.to_string());
        let vector_search = self.vector_search.clone();
        
        let conn = self.connection.clone();
        let result = tokio::task::spawn_blocking(move || {
            let conn = conn.lock().unwrap();
            
            // First update the database record
            let mut stmt = conn.prepare(
                "UPDATE knowledge_items SET title = COALESCE(?, title), content = COALESCE(?, content), 
                 tags = COALESCE(?, tags), metadata = COALESCE(?, metadata), updated_at = CURRENT_TIMESTAMP 
                 WHERE id = ?"
            )?;
            
            let rows_affected = stmt.execute((title.as_deref(), content.as_deref(), tags.as_deref(), metadata.as_deref(), id))?;
            
            if rows_affected > 0 {
                // Get the updated item for vector re-indexing
                let mut stmt = conn.prepare(
                    "SELECT id, uuid, title, content, item_type, url, file_path, file_size, mime_type, 
                     created_at, updated_at, tags, metadata FROM knowledge_items WHERE id = ?"
                )?;
                
                let item = stmt.query_row([id], |row| {
                    Ok(KnowledgeItem {
                        id: row.get(0)?,
                        uuid: row.get(1)?,
                        title: row.get(2)?,
                        content: row.get(3)?,
                        item_type: row.get(4)?,
                        url: row.get(5)?,
                        file_path: row.get(6)?,
                        file_size: row.get(7)?,
                        mime_type: row.get(8)?,
                        created_at: row.get(9)?,
                        updated_at: row.get(10)?,
                        tags: row.get(11)?,
                        metadata: row.get(12)?,
                    })
                })?;
                
                Ok::<(bool, KnowledgeItem), anyhow::Error>((true, item))
            } else {
                Ok::<(bool, KnowledgeItem), anyhow::Error>((false, KnowledgeItem::default()))
            }
        }).await??;
        
        if result.0 {
            // Update vector index if vector search is available
            if let Some(vector_search) = &vector_search {
                if let Err(e) = vector_search.update_knowledge_item_with_agent(&result.1, agent_id).await {
                    tracing::error!("Failed to update vector index: {}", e);
                }
            }
            Ok(true)
        } else {
            Ok(false)
        }
    }

    pub async fn delete_knowledge_item_with_vector(&self, id: i64) -> Result<bool> {
        // Remove from vector store first
        if let Some(vector_search) = &self.vector_search {
            if let Err(e) = vector_search.remove_knowledge_item(id).await {
                tracing::warn!("Failed to remove item from vector store: {}", e);
            }
        }

        // Delete from database
        let conn = self.connection.lock().unwrap();
        let rows_affected = conn.execute("DELETE FROM knowledge_items WHERE id = ?1", [id])?;
        Ok(rows_affected > 0)
    }

    pub async fn batch_add_knowledge_items_with_vector(
        &self,
        items: Vec<(String, Option<String>, String)>, // (title, content, item_type)
    ) -> Result<Vec<i64>> {
        self.batch_add_knowledge_items_with_vector_and_agent(items, None).await
    }

    pub async fn batch_add_knowledge_items_with_vector_and_agent(
        &self,
        items: Vec<(String, Option<String>, String)>, // (title, content, item_type)
        agent_id: Option<i64>,
    ) -> Result<Vec<i64>> {
        let mut item_ids = Vec::new();
        let mut created_items = Vec::new();

        // Create items in database
        for (title, content, item_type) in items {
            let item_id = self.create_knowledge_item(
                &title,
                content.as_deref(),
                &item_type,
                None,
                None,
                None,
                None,
                None,
                None,
            )?;
            item_ids.push(item_id);

            if let Ok(Some(item)) = self.get_knowledge_item(item_id) {
                created_items.push(item);
            }
        }

        // Batch add to vector store with agent configuration
        if let Some(vector_search) = &self.vector_search {
            let item_refs: Vec<&KnowledgeItem> = created_items.iter().collect();
            if let Err(e) = vector_search.batch_add_knowledge_items_with_agent(item_refs, agent_id).await {
                tracing::warn!("Failed to batch add items to vector store: {}", e);
            }
        }

        Ok(item_ids)
    }
}
